/**
 * Copyright (c) 2023 myons Inc. All rights reserved.
 */
package cn.zmwh.common.security.base;

import cn.hutool.core.util.StrUtil;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

/**
 * @describe: 修改header 、
 * @author: dmzmwh 、
 * @time: 2018-03-10 15:01
 */
public class MutableHttpServletRequest extends HttpServletRequestWrapper {
    private final Map<String, String> headerMap = new HashMap<>();

    private Map<String, String[]> paramsSubjoin = new HashMap<>();
    // 重新赋值的body数据
    private String bodyJsonStr;

    public MutableHttpServletRequest(HttpServletRequest request) {
        super(request);
    }

    public void putHeader(String name, String value) {
        this.headerMap.put(name, value);
    }

    public void putHeaderMap(Map<String, String> map) {
        this.headerMap.putAll(map);
    }

    @Override
    public String getHeader(String name) {
        String headerValue = headerMap.get(name);

        if (headerValue != null) {
            return headerValue;
        }
        return super.getHeader(name);
    }

    @Override
    public Enumeration<String> getHeaderNames() {
        if (headerMap.isEmpty()) {
            return super.getHeaderNames();
        }

        Set<String> set = new HashSet<String>(headerMap.keySet());
        // 添加自定义header
        Enumeration<String> e = super.getHeaderNames();
        while (e.hasMoreElements()) {
            String n = e.nextElement();
            set.add(n);
        }

        return Collections.enumeration(set);
    }

    @Override
    public Enumeration<String> getHeaders(String name) {
        List<String> values = Collections.list(super.getHeaders(name));
        if (this.headerMap.containsKey(name)) {
            values = Arrays.asList(this.headerMap.get(name));
        }
        return Collections.enumeration(values);
    }


    @Override
    public Enumeration<String> getParameterNames() {
        if (paramsSubjoin.isEmpty()) {
            return super.getParameterNames();
        }

        Set<String> set = new HashSet<String>(paramsSubjoin.keySet());
        // 添加自定义header
        Enumeration<String> e = super.getParameterNames();
        while (e.hasMoreElements()) {
            String n = e.nextElement();
            set.add(n);
        }

        return Collections.enumeration(set);
    }

    @Override
    public String getParameter(String name) {
        String[] values = paramsSubjoin.get(name);
        if (values == null || values.length == 0) {
            return super.getParameter(name);
        }
        return values[0];
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] values = paramsSubjoin.get(name);
        if (values == null || values.length == 0) {
            return super.getParameterValues(name);
        }
        return values;
    }

    private void addAllParameters(Map<String, Object> otherParams) {
        for (Map.Entry<String, Object> entry : otherParams.entrySet()) {
            addParameter(entry.getKey(), entry.getValue());
        }
    }

    public void addParameter(String name, Object value) {
        if (value != null) {
            if (value instanceof String[]) {
                paramsSubjoin.put(name, (String[]) value);
            } else if (value instanceof String) {
                paramsSubjoin.put(name, new String[]{(String) value});
            } else {
                paramsSubjoin.put(name, new String[]{String.valueOf(value)});
            }
        }
    }


    @Override
    public ServletInputStream getInputStream() throws IOException {
        if (StrUtil.isBlank(bodyJsonStr)) {
            bodyJsonStr = "";
        }
        // 必须指定utf-8编码，否则json请求数据中如果包含中文，会出现异常
        final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(bodyJsonStr.getBytes("utf-8"));
        ServletInputStream servletInputStream = new ServletInputStream() {
            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener readListener) {
            }

            @Override
            public int read() throws IOException {
                return byteArrayInputStream.read();
            }
        };
        return servletInputStream;
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(this.getInputStream()));
    }

    public String getBodyJsonStr() {
        return bodyJsonStr;
    }

    public void setBodyJsonStr(String bodyJsonStr) {
        this.bodyJsonStr = bodyJsonStr;
    }

}